/**
 * PHOSIDE: PHosphorylation Site IDentification Engine.
 * Copyright 2009 Digital Biology Lab, University of Missouri.
 * This library is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation; either version 3 of the License, or (at your option) any later
 * version. <p/> This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the License for more
 * details.
 */

package phoside.io.xml;

import java.io.File;
import java.io.IOException;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang.StringUtils;

import org.dom4j.Element;
import org.dom4j.ElementPath;
import org.dom4j.ElementHandler;
import org.dom4j.DocumentException;
import org.dom4j.io.SAXReader;

import phoside.PhosphoProteins;
import phoside.ProteinImpl;

import phoside.io.ProteinsReader;

/**
 *
 * @author gjj
 */
public class UniProtPhosphoReader implements ProteinsReader {
    protected final String dirXml;
    protected final Set<PhosphoProteins.PhosphoType> types;
    protected final Set<String> organism;
    protected final PhosphoProteins data;
    protected final boolean includeBySimilarity;

    private boolean keepAllIds = false;
    private Map<String,String> mapIdMainId = null;

    public UniProtPhosphoReader(String dirXml, PhosphoProteins data) {
        this (dirXml, data, null);
    }

    public UniProtPhosphoReader(String dirXml, PhosphoProteins data,
            Set<PhosphoProteins.PhosphoType> types) {
        this (dirXml, data, types, null);
    }

    public UniProtPhosphoReader(String dirXml, PhosphoProteins data,
            Set<PhosphoProteins.PhosphoType> types, Set<String> organism) {
        this (dirXml, data, types, organism, true);
    }

    /**
     *
     * @param dirXml
     * @param types
     * @param organism if organism is null, read all
     * @param data
     * @param includeBySimilarity
     */
    public UniProtPhosphoReader(String dirXml, PhosphoProteins data,
            Set<PhosphoProteins.PhosphoType> types,
            Set<String> organism, boolean includeBySimilarity) {
        this.dirXml = dirXml;
        this.types = types;
        this.organism = organism;
        this.data = data;
        this.includeBySimilarity = includeBySimilarity;
    }

    public void setKeepAllIds() {
        this.keepAllIds = true;
        mapIdMainId = new HashMap();
    }

    public Map getIDMap() {
        return mapIdMainId;
    }

    public PhosphoProteins read() throws IOException {
        File file = new File(dirXml);
        if (!file.exists()) {
            throw new java.io.FileNotFoundException(dirXml+" cannot be found.");
        }

        final Set<Character> allowTypes;
        if (types==null) {
            allowTypes = PhosphoProteins.PhosphoType.phosphorylatableAminoAcids();
        } else {
            allowTypes = PhosphoProteins.PhosphoType.phosphorylatableAminoAcids(types);
        }

        SAXReader reader = new SAXReader();

        final StringBuilder acc = new StringBuilder(30);
        final StringBuilder name = new StringBuilder(30);
        final StringBuilder fullName = new StringBuilder(200);
        final StringBuilder org = new StringBuilder(30);
        final StringBuilder seq = new StringBuilder(2000);
        final Map<Integer, Set<String>> sites = new HashMap();
        final Map<Integer, Boolean> bySim = new HashMap();
        final Set<String> accs = new HashSet();

        // entry
        reader.addHandler( "/uniprot/entry", new ElementHandler() {
            public void onStart(ElementPath path) {
                acc.setLength(0);
                fullName.setLength(0);
                seq.setLength(0);
                org.setLength(0);
                name.setLength(0);
                sites.clear();
                bySim.clear();
                accs.clear();
            }
            public void onEnd(ElementPath path) {
                // process a element
                if (org.length()>0 && (organism==null || 
                        organism.contains(org.toString()))
                        && acc.length()>0 && seq.length()>0) {
                    String accession = acc.toString();
                    String sequence = seq.toString();
                    data.addProtein(new ProteinImpl(acc.toString(),
                            sequence, name.toString(), fullName.toString(),
                            org.toString()));

                    for (int site : sites.keySet()) {
                        if (site<0 || site>=sequence.length())
                            continue;
                        if (!allowTypes.contains(sequence.charAt(site)))
                            continue;

                        List<String> notes = new ArrayList();

                        if (bySim.get(site)) {
                            notes.add("by_similarity");
                        }

                        Set<String> kinases = sites.get(site);
                        if (kinases!=null && !kinases.isEmpty()) {
                            notes.add("kinase:"+StringUtils.join(kinases, ','));
                        }

                        data.addSite(accession, site, notes.isEmpty()?null:notes);
                    }

                    if (keepAllIds) {
                        for (String ac : accs) {
                            mapIdMainId.put(ac, accession);
                        }
                    }

                    //System.out.println(accession);
                }

                // prune the tree
                Element row = path.getCurrent();
                row.detach();
            }
        });

        // accession
        reader.addHandler("/uniprot/entry/accession",
                new ElementHandler() {
            public void onStart(ElementPath path) {
                // do nothing
            }

            public void onEnd(ElementPath path) {
                if (acc.length()==0) {
                    Element el = path.getCurrent();
                    acc.append(el.getText());
                    if (keepAllIds) {
                        accs.add(acc.toString());
                    }
                } else {
                    if (keepAllIds) {
                        accs.add(path.getCurrent().getText());
                    }
                }
                
            }
        });

        // name
        reader.addHandler("/uniprot/entry/name",
                new ElementHandler() {
            public void onStart(ElementPath path) {
                // do nothing
            }

            public void onEnd(ElementPath path) {
                if (name.length()>0)
                    return;

                Element el = path.getCurrent();
                name.append(el.getText());
            }
        });

        // full name
        reader.addHandler("/uniprot/entry/protein/recommendedName/fullName",
                new ElementHandler() {
            public void onStart(ElementPath path) {
                // do nothing
            }

            public void onEnd(ElementPath path) {
                if (fullName.length()>0)
                    return;

                Element el = path.getCurrent();
                fullName.append(el.getTextTrim());
            }
        });

        reader.addHandler("/uniprot/entry/organism/name",
                new ElementHandler() {
            public void onStart(ElementPath path) {
                // do nothing
            }

            public void onEnd(ElementPath path) {
                if (org.length()>0)
                    return;

                Element el = path.getCurrent();
                String attr = el.attributeValue("type");
                if (attr==null || !attr.equalsIgnoreCase("scientific")) {
                    return;
                }

                org.append(el.getText());
            }
        });

        reader.addHandler("/uniprot/entry/sequence",
                new ElementHandler() {
            public void onStart(ElementPath path) {
                // do nothing
            }

            public void onEnd(ElementPath path) {
                if (seq.length()>0)
                    return;

                Element el = path.getCurrent();
                seq.append(el.getText().replaceAll("\\p{Space}", ""));
            }
        });

        reader.addHandler("/uniprot/entry/feature",
                new ElementHandler() {
            public void onStart(ElementPath path) {
                // do nothing
            }

            public void onEnd(ElementPath path) {
                Element el = path.getCurrent();
                String attr = el.attributeValue("type");
                if (attr==null || !attr.equalsIgnoreCase("modified residue")) 
                    return;

                attr = el.attributeValue("description");
                if (attr==null || !attr.startsWith("Phospho"))
                    return;

                String kinase = null;
                int idx = attr.indexOf("; by ");
                if (idx!=-1) {
                    kinase = attr.substring(idx+5);
                    idx = kinase.indexOf(" ");
                    if (idx!=-1) {
                        kinase = kinase.substring(0, idx);
                    }
                }

                attr = el.attributeValue("status");
                boolean sim = attr!=null && attr.toLowerCase().contains("by similarity");
                if (!includeBySimilarity && sim) {
                    return;
                }

                List<Element> locs = el.elements("location");
                for (Element loc : locs) {
                    List<Element> poss = loc.elements("position");
                    for (Element pos : poss) {
                        String str = pos.attributeValue("position");
                        if (str==null) continue;

                        try {
                            int site = Integer.parseInt(str)-1; //start from 0
                            //sites.add(site);
                            Set<String> kinases = sites.get(site);
                            if (kinases==null) {
                                kinases = new HashSet();
                                sites.put(site, kinases);
                            }
                            if (kinase!=null)
                                kinases.add(kinase);
                            bySim.put(site,sim);
                        } catch (NumberFormatException e) {
                            continue;
                        }
                    }
                }

            }
        });
        

        try {
            reader.read(file);
        } catch (DocumentException e) {
            e.printStackTrace();
        }        

        return data;
    }
}
